﻿using System;
using System.Collections.Generic;
using System.Linq;
using Microsoft.Build.Utilities;
using Microsoft.Build.Framework;
using System.IO;
using System.Security.Cryptography;

namespace TfsCloudDeploy.Tasks
{
    public class GetModifiedAssemblies : Task
    {
        #region Constants

        private const string TEMP_TOOL_FOLDER = "TempTools";
        private const string TOOL_NAME = "PEedit.exe";

        #endregion

        #region Private Fields

        private IEnumerable<ITaskItem> _builtBinaryFiles;
        private IEnumerable<ITaskItem> _builtDebugSymbolFiles;
        private IEnumerable<ITaskItem> _builtXmlDocumentationFiles;
        private string _toolPath;
        private string _tempFolder;

        #endregion

        #region Constructors

        public GetModifiedAssemblies() { }

        #endregion

        #region Task Properties

        [Required]
        public string BaselineAssembliesFolder { get; set; }

        [Required]
        public ITaskItem[] BuiltAssemblies { get; set; }

        /// <summary>
        /// The full path to the tool including the tool name.
        /// </summary>
        public string ToolPath { get; set; }

        public bool SkipRemoveTempFolder { get; set; }

        public bool ShowVerboseLogOutput { get; set; }

        public bool IgnoreAssemblyVersionChange { get; set; }

        public bool IgnoreAssemblyFileVersionChange { get; set; }

        [Output]
        public ITaskItem[] ModifiedAssemblies { get; private set; }

        #endregion

        #region Execute

        public override bool Execute()
        {
            try
            {
                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Started at {0:MM/dd/yyyy H:mm:ss:fff zzz}.", DateTime.Now);

                // Are there any assemblies provided for comparison?
                if (BuiltAssemblies.Length == 0)
                {
                    Log.LogError("GetModifiedAssemblies: BuiltAssemblies list is empty.");

                    return false;
                }

                // Get the tool path:
                _toolPath = this.getToolPath();
                if (!File.Exists(_toolPath))
                {
                    Log.LogWarning("GetModifiedAssemblies: Tool not found at path: {0}.", _toolPath);
                    Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Skipping assembly comparison.");

                    // Return the original list of assemblies since we cannot compare them.
                    ModifiedAssemblies = BuiltAssemblies;

                    return true;
                }

                // Is a baseline directory provided?
                if (string.IsNullOrWhiteSpace(BaselineAssembliesFolder) || !Directory.Exists(BaselineAssembliesFolder))
                {
                    Log.LogWarning("GetModifiedAssemblies: BaselineAssembliesFolder is empty or folder does not exist: {0}.", BaselineAssembliesFolder);
                    Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Skipping assembly comparison.");

                    // Return the original list of assemblies since we cannot compare them.
                    ModifiedAssemblies = BuiltAssemblies;

                    return true;
                }

                _tempFolder = createTemporaryWorkingFolder();

                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Created temp folder at: {0}.", _tempFolder);
                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: BaselineAssembliesFolder: {0}.", BaselineAssembliesFolder);
                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: SkipRemoveTempFolder: {0}.", SkipRemoveTempFolder);
                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: BuiltAssemblies count: {0}.", BuiltAssemblies.Count());

                // break down the input into 3 groups:
                _builtXmlDocumentationFiles = from file in BuiltAssemblies
                                              let lowercaseFileName = file.ItemSpec.ToLowerInvariant()
                                              where lowercaseFileName.EndsWith(".xml")
                                              select file;

                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: XML files count: {0}.", _builtXmlDocumentationFiles.Count());

                _builtDebugSymbolFiles = from file in BuiltAssemblies
                                         let lowercaseFileName = file.ItemSpec.ToLowerInvariant()
                                         where lowercaseFileName.EndsWith(".pdb")
                                         select file;

                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: PDB files count: {0}.", _builtDebugSymbolFiles.Count());

                _builtBinaryFiles = from file in BuiltAssemblies
                                    let lowercaseFileName = file.ItemSpec.ToLowerInvariant()
                                    where lowercaseFileName.EndsWith(".dll") || lowercaseFileName.EndsWith(".exe")
                                    select file;

                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Binary files count: {0}.", _builtBinaryFiles.Count());

                int counter = 0;
                List<ITaskItem> modifiedAssemblies = new List<ITaskItem>();

                // iterate through the DLL/EXE files and compare them to the baseline:
                foreach (var item in _builtBinaryFiles)
                {
                    if (!isAssemblyUnchanged(item.ItemSpec))
                    {
                        Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Found changed assembly: {0}.", item.ItemSpec);

                        // Add .dll/.exe
                        modifiedAssemblies.Add(new TaskItem(item));
                        counter++;

                        // any .pdb?
                        var pdbFile = _builtDebugSymbolFiles.FirstOrDefault(file => string.Compare(file.ItemSpec, string.Format("{0}.pdb", item.ItemSpec.Substring(0, item.ItemSpec.Length - 4)), StringComparison.InvariantCultureIgnoreCase) == 0);
                        if (pdbFile != null)
                        {
                            Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Adding PDB file: {0}.", pdbFile.ItemSpec);
                            modifiedAssemblies.Add(pdbFile);
                        }

                        // any .xml?
                        var xmlFile = _builtXmlDocumentationFiles.FirstOrDefault(file => string.Compare(file.ItemSpec, string.Format("{0}.xml", item.ItemSpec.Substring(0, item.ItemSpec.Length - 4)), StringComparison.InvariantCultureIgnoreCase) == 0);
                        if (xmlFile != null)
                        {
                            Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Adding XML file: {0}.", xmlFile.ItemSpec);
                            modifiedAssemblies.Add(xmlFile);
                        }
                    }
                    else
                    {
                        Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Skipping unmodified assembly: {0}.", item.ItemSpec);
                    }
                }

                ModifiedAssemblies = modifiedAssemblies.ToArray();

                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: {0} modified assemblies found ({1} total artifacts).", counter, ModifiedAssemblies.Length);

                return true;
            }
            catch (Exception ex)
            {
                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Exception during task execution.");
                Log.LogErrorFromException(ex);
                return false;
            }
            finally
            {
                cleanTempFolder();

                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Finished at {0:MM/dd/yyyy H:mm:ss:fff zzz}.", DateTime.Now);
            }
        }

        private string getToolPath()
        {
            if (!string.IsNullOrWhiteSpace(ToolPath))
            {
                return ToolPath;
            }

            // If there is a cached tool in the temp tool path, use it:
            string toolPath = getTemporaryToolPath();
            if (!string.IsNullOrWhiteSpace(toolPath) && File.Exists(toolPath))
            {
                return toolPath;
            }

            // If there is an embedded image, extract to the temp location.
            byte[] buffer = Properties.Resources.peedit;
            if (buffer != null && buffer.Length > 0)
            {
                FileInfo fi = new FileInfo(toolPath);
                DirectoryInfo di = fi.Directory;
                if (!di.Exists)
                {
                    di.Create();
                }

                using (FileStream fs = File.Create(toolPath))
                {
                    fs.Write(buffer, 0, buffer.Length);
                }
                return toolPath;
            }

            return PEHeaderUtil.DefaultToolPath;
        }

        private void cleanTempFolder()
        {
            Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: cleanTempFolder [SkipRemoveTempFolder={0}]: {1}.", SkipRemoveTempFolder, _tempFolder);

            // Cleanup the temporary folder if one was created:
            if (!SkipRemoveTempFolder && !string.IsNullOrWhiteSpace(_tempFolder))
            {
                DirectoryInfo tempFolderInfo = new DirectoryInfo(_tempFolder);
                try
                {
                    if (tempFolderInfo.Exists)
                    {
                        tempFolderInfo.Delete(true);
                    }
                }
                catch (IOException ex)
                {
                    Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: IOException in cleanTempFolder.");
                    Log.LogWarningFromException(ex);
                    Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Failed while cleaning up temp folder: {0}.", _tempFolder);

                }
            }
        }

        private bool isAssemblyUnchanged(string filename)
        {
            try
            {
                FileInfo fileInfo = new FileInfo(filename);
                if (fileInfo.Exists)
                {
                    Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Processing file: {0}.", filename);

                    // get the baseline file:
                    string baselineFile = getFilePathUnderRootFolder(fileInfo.Name, BaselineAssembliesFolder);
                    if (!string.IsNullOrWhiteSpace(baselineFile))
                    {
                        Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Corresponding baseline file: {0}.", baselineFile);

                        SimpleHeader baselineHeader = PEHeaderUtil.GetSimpleHeader(baselineFile, _toolPath);

                        if (baselineHeader != null)
                        {
                            if (ShowVerboseLogOutput)
                            {
                                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: baselineHeader: {0}.",
                                               baselineHeader.ToString());
                            }

                            string baselineFileHash = ComputeHash(baselineFile);
                            string builtFileHash = getBuiltFileHash(fileInfo, baselineHeader);

                            if (builtFileHash != null)
                            {
                                if (ShowVerboseLogOutput)
                                {
                                    Log.LogMessage(MessageImportance.High,
                                                   "GetModifiedAssemblies: baselineFileHash: {0}.", baselineFileHash);
                                    Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: builtFileHash: {0}.",
                                                   builtFileHash);
                                }

                                // Assembly is unchanged if both hash values are equal.
                                return (string.Compare(baselineFileHash, builtFileHash, StringComparison.InvariantCultureIgnoreCase) == 0);
                            }
                        }
                        else
                        {
                            Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: baselineHeader is null.");
                        }
                    }
                    else
                    {
                        // This could be due to this being the first time a file is being checked in, so it's not
                        // present yet in the baseline folder.
                        Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Baseline file not found.");
                    }
                }
                else
                {
                    Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: File not found: {0}.", filename);
                }

                return false;
            }
            catch (Exception ex)
            {
                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Exception in isAssemblyUnchanged.");
                Log.LogErrorFromException(ex);
                return false;
            }
        }

        private string getFilePathUnderRootFolder(string filename, string rootFolderPath)
        {
            if (string.IsNullOrWhiteSpace(filename) || !Directory.Exists(rootFolderPath))
            {
                return null;
            }

            string[] files = Directory.GetFiles(rootFolderPath, filename, SearchOption.AllDirectories);
            if (files.Length > 0)
            {
                return files[0];
            }

            return null;
        }

        private string getBuiltFileHash(FileInfo builtFile, SimpleHeader baselineHeader)
        {
            try
            {
                // Copy the built file to %TEMP%.
                string builtFilePath = Path.Combine(_tempFolder, builtFile.Name);
                File.Copy(builtFile.FullName, builtFilePath, true);
                FileInfo builtFileInfo = new FileInfo(builtFilePath);

                if (builtFileInfo.Exists)
                {
                    builtFileInfo.IsReadOnly = false; // remove read-only flag.

                    // If DebugDirectoryEntryMvid is blank, then set it to the FileHeaderMvid.
                    // This is typically the case where an assembly was built without any debug information
                    // generated (i.e. no PDB is created, so no need to have a debug entry that points to it).
                    if (string.IsNullOrWhiteSpace(baselineHeader.DebugDirectoryEntryMvid))
                        baselineHeader.DebugDirectoryEntryMvid = baselineHeader.FileHeaderMvid;

                    // Modify its header.
                    SimpleHeader modifiedHeader = (!IgnoreAssemblyVersionChange || !IgnoreAssemblyFileVersionChange) ?
                        PEHeaderUtil.SetSimpleHeaderWithoutVersionChange(baselineHeader, builtFilePath, _toolPath, !IgnoreAssemblyFileVersionChange, !IgnoreAssemblyVersionChange) :
                        PEHeaderUtil.SetSimpleHeader(baselineHeader, builtFilePath, _toolPath);

                    if (modifiedHeader != null)
                    {

                        if (ShowVerboseLogOutput)
                        {
                            Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: modifiedHeader: {0}.", modifiedHeader.ToString());
                        }

                        // Compute the hash.
                        return ComputeHash(builtFilePath);
                    }
                    else
                    {
                        Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: modifiedHeader is null.");
                        return null;
                    }
                }
                else
                {
                    Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: getBuiltFileHash: File not found: {0}.", builtFilePath);
                    return null;
                }
            }
            catch (Exception ex)
            {
                Log.LogMessage(MessageImportance.High, "GetModifiedAssemblies: Exception in getBuiltFileHash.");
                Log.LogErrorFromException(ex);
                return null;
            }

        }

        #endregion

        #region Static Members

        private static string createTemporaryWorkingFolder()
        {
            string tempFolder = Environment.GetEnvironmentVariable("TEMP");
            string workingFolderName = string.Format("PE_{0}", new Random().Next(10000000, 99999999));
            DirectoryInfo directory = new DirectoryInfo(tempFolder);
            DirectoryInfo subDirectory = directory.CreateSubdirectory(workingFolderName);
            return subDirectory.FullName;
        }

        private static string getTemporaryToolPath()
        {
            string tempFolder = Environment.GetEnvironmentVariable("TEMP");
            return Path.Combine(tempFolder, TEMP_TOOL_FOLDER, TOOL_NAME);
        }

        public static string ComputeHash(string fileName)
        {
            SHA1 hashAlgorithm = SHA1CryptoServiceProvider.Create();
            using (FileStream stream = File.OpenRead(fileName))
            {
                byte[] hash = hashAlgorithm.ComputeHash(stream);
                return Convert.ToBase64String(hash);
            }
        }

        #endregion

    }
}
